# -*- coding: utf-8 -*-
# time: 2025/5/12 16:08
# file: ch04_ds_sql.py
# author: hanson
"""
连接数据库，读操作/数据查询等。
1, 用户输入     转sql查询。查询的结果给大模型 组织结果，--》最后返回给用户。
               不转sql查询。  直接大模型给结果       --》最后返回给用户。

sqlalchemy 是langchain内置的链接数据库。
"""
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.agent_toolkits import create_sql_agent, SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_core.tools import Tool
from langchain_ollama import ChatOllama
from sqlalchemy import create_engine

# engine = create_engine("mysql+mysqldb://root:1234qwer@192.168.6.43:3306/test", pool_recycle=3)
engine = create_engine("mysql+mysqldb://root:1234qwer@192.168.6.43:3306/iot_ts", pool_recycle=3600)
# 创建SQLDatabase实例 可以限制表 include_tables=['users','orders'
db = SQLDatabase(engine)


# 测试 列出表名
print(db.get_table_names())
# 测试查询
print(db.run("SELECT * FROM users WHERE age > 35"))

# 构建查询工具
llm = ChatOllama(model="qwen2.5:1.5b", temperature=0.3)


# 构建Toolkit
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
# 创建SQL Agent
agent = create_sql_agent(
    llm=llm,  # 需替换为实际LLM实例
    toolkit=toolkit,
    verbose=True
)
# 大模型 只能根据你的问题，生成sql语句，并不会执行sql语句 ，如果换成"qwen2.5:7b" 可以查出结果。
result = agent.run("查询年龄大于35的用户(users)的name值")
print(result)

query_chain = create_sql_query_chain(llm, db)
result = query_chain.invoke({'question':"查询年龄大于35的用户的名称"})
print(result)


